#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MNIST + 用户手写数字识别（Tkinter GUI）
功能：
1. 首次运行自动训练并保存模型；
2. 以后直接加载模型，秒进 GUI；
3. 菜单栏“模型→重新训练”可强制从头训练并覆盖旧模型；
4. 原“重新训练模型”按钮仍保留，仅供快速微调。
"""

import os
import pickle
import numpy as np
from PIL import Image, ImageDraw
import tkinter as tk
from tkinter import simpledialog, messagebox

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset

# ----------- 路径常量 -----------
DATA_DIR = './data'
USER_DATA_FILE = os.path.join(DATA_DIR, 'user_handwritten_data.pkl')
MODEL_FILE = os.path.join(DATA_DIR, 'simple_cnn.pth')
os.makedirs(DATA_DIR, exist_ok=True)

# ----------- 数据预处理 -----------
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# ----------- 用户数据集 -----------
class UserDataset(Dataset):
    def __init__(self, data=None, transform=None):
        self.data = data if data is not None else []
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        image, label = self.data[idx]
        if self.transform:
            image = self.transform(image)
        return image, label

    def add_sample(self, image, label):
        image_28 = image.resize((28, 28), Image.LANCZOS).convert('L')
        self.data.append((image_28, label))

    def save_to_file(self, filename):
        with open(filename, 'wb') as f:
            pickle.dump(self.data, f)

    def load_from_file(self, filename):
        if os.path.exists(filename):
            with open(filename, 'rb') as f:
                self.data = pickle.load(f)

# ----------- 加载 MNIST -----------
trainset = torchvision.datasets.MNIST(root=DATA_DIR, train=True, download=True, transform=transform)
testset = torchvision.datasets.MNIST(root=DATA_DIR, train=False, download=True, transform=transform)

user_dataset = UserDataset(transform=transform)
user_dataset.load_from_file(USER_DATA_FILE)
combined_dataset = torch.utils.data.ConcatDataset([trainset, user_dataset])

trainloader = DataLoader(combined_dataset, batch_size=64, shuffle=True)
testloader = DataLoader(testset, batch_size=64, shuffle=False)

# ----------- 模型定义 -----------
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.fc1 = nn.Linear(20 * 4 * 4, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1, 20 * 4 * 4)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# ----------- 工具函数 -----------
def train_model(model, loader, criterion, optimizer, epochs):
    model.train()
    for epoch in range(epochs):
        running_loss = 0.0
        for i, (inputs, labels) in enumerate(loader, 0):
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
            if i % 200 == 199:
                print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 200:.3f}')
                running_loss = 0.0

def test_model(model, loader):
    model.eval()
    correct = total = 0
    with torch.no_grad():
        for inputs, labels in loader:
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    acc = 100 * correct / total
    print(f'Accuracy on 10000 test images: {acc:.2f}%')
    return acc

def save_model(model, path=MODEL_FILE):
    torch.save(model.state_dict(), path)
    print(f'[INFO] 模型已保存到 {path}')

def load_model(model, path=MODEL_FILE):
    if os.path.isfile(path):
        model.load_state_dict(torch.load(path, map_location='cpu'))
        print(f'[INFO] 模型已加载自 {path}')
        return True
    return False

# ----------- 初始模型加载或训练 -----------
model = SimpleCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)

if load_model(model):
    initial_accuracy = test_model(model, testloader)
else:
    print('首次运行，开始训练模型...')
    train_model(model, trainloader, criterion, optimizer, epochs=5)
    initial_accuracy = test_model(model, testloader)
    save_model(model)

# ----------- GUI 构建 -----------
root = tk.Tk()
root.title('手写数字识别')

# 菜单栏
menubar = tk.Menu(root)
root.config(menu=menubar)
model_menu = tk.Menu(menubar, tearoff=0)
menubar.add_cascade(label='模型', menu=model_menu)

# 准确率标签
accuracy_label = tk.Label(root, text=f'模型准确率: {initial_accuracy:.2f}%')
accuracy_label.pack()

# 画布
canvas_width = canvas_height = 280
canvas = tk.Canvas(root, width=canvas_width, height=canvas_height, bg='white')
canvas.pack()

# 画表格
for i in range(1, 28):
    canvas.create_line(i * 10, 0, i * 10, canvas_height, fill='gray')
    canvas.create_line(0, i * 10, canvas_width, i * 10, fill='gray')

# PIL 图像
pil_image = Image.new('L', (280, 280), 0)
draw = ImageDraw.Draw(pil_image)

def paint(event):
    x1, y1 = event.x - 5, event.y - 5
    x2, y2 = event.x + 5, event.y + 5
    canvas.create_oval(x1, y1, x2, y2, fill='black', width=10)
    draw.ellipse([x1, y1, x2, y2], fill=255, width=10)

canvas.bind('<B1-Motion>', paint)

# ----------- 识别与纠正 -----------
def recognize():
    resized = pil_image.resize((28, 28), Image.LANCZOS).convert('L')
    input_tensor = transform(resized).unsqueeze(0)
    model.eval()
    with torch.no_grad():
        output = model(input_tensor)
        prob = torch.nn.functional.softmax(output, dim=1)[0]
        confidence, predicted = torch.max(prob, 0)
        confidence, predicted = confidence.item() * 100, predicted.item()

    result_label.config(text=f'预测结果: {predicted} (置信度: {confidence:.2f}%)')

    if confidence < 70 or not messagebox.askyesno('确认', f'系统识别为 {predicted}，是否正确？'):
        correct_and_add(predicted, pil_image)
    else:
        if messagebox.askyesno('添加训练数据', '是否将此样本添加到训练集以改进模型？'):
            user_dataset.add_sample(pil_image, predicted)
            user_dataset.save_to_file(USER_DATA_FILE)
            messagebox.showinfo('成功', f'已将数字 {predicted} 添加到训练集！')

def correct_and_add(recognized_label, user_image):
    correct_label = simpledialog.askinteger('纠正识别结果',
                                            f'系统识别为: {recognized_label}\n请输入正确数字 (0-9):',minvalue=0, maxvalue=9)
    if correct_label is not None:
        user_dataset.add_sample(user_image, correct_label)
        user_dataset.save_to_file(USER_DATA_FILE)
        global combined_dataset, trainloader
        combined_dataset = torch.utils.data.ConcatDataset([trainset, user_dataset])
        trainloader = DataLoader(combined_dataset, batch_size=64, shuffle=True)
        messagebox.showinfo('成功', f'已将数字 {correct_label} 添加到训练集！')

# ----------- 微调模型 -----------
def retrain_model():
    if len(user_dataset) == 0:
        messagebox.showinfo('提示', '没有新的用户数据需要训练！')
        return
    messagebox.showinfo('开始训练', f'使用 {len(user_dataset)} 个用户样本进行微调...')
    train_model(model, trainloader, criterion, optimizer, epochs=2)
    acc = test_model(model, testloader)
    accuracy_label.config(text=f'模型准确率: {acc:.2f}%')
    save_model(model)
    messagebox.showinfo('训练完成', f'微调完成！新准确率: {acc:.2f}%')

# ----------- 菜单栏：从头训练 -----------
def retrain_model_from_scratch():
    ans = messagebox.askyesno('确认', '将从头重新训练模型并覆盖旧模型，是否继续？')
    if not ans:
        return
    global model, optimizer
    model = SimpleCNN()
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
    print('重新训练中...')
    train_model(model, trainloader, criterion, optimizer, epochs=5)
    acc = test_model(model, testloader)
    save_model(model)
    accuracy_label.config(text=f'模型准确率: {acc:.2f}%')
    messagebox.showinfo('完成', f'重新训练结束，准确率 {acc:.2f}%')

model_menu.add_command(label='重新训练（覆盖旧模型）', command=retrain_model_from_scratch)

# ----------- 清除画布 -----------
def clear():
    canvas.delete('all')
    draw.rectangle([0, 0, 280, 280], fill=0)
    result_label.config(text='预测结果: ')
    for i in range(1, 28):
        canvas.create_line(i * 10, 0, i * 10, canvas_height, fill='gray')
        canvas.create_line(0, i * 10, canvas_width, i * 10, fill='gray')

# ----------- 按钮与标签 -----------
button_frame = tk.Frame(root)
button_frame.pack(pady=5)

recognize_button = tk.Button(button_frame, text='识别', command=recognize)
recognize_button.pack(side=tk.LEFT, padx=5)

retrain_button = tk.Button(button_frame, text='重新训练模型（微调）', command=retrain_model)
retrain_button.pack(side=tk.LEFT, padx=5)

clear_button = tk.Button(button_frame, text='清除', command=clear)
clear_button.pack(side=tk.LEFT, padx=5)

result_label = tk.Label(root, text='预测结果: ')
result_label.pack(pady=5)

data_count_label = tk.Label(root, text=f'已收集用户样本数: {len(user_dataset)}')
data_count_label.pack(pady=5)

def update_data_count():
    data_count_label.config(text=f'已收集用户样本数: {len(user_dataset)}')
    root.after(1000, update_data_count)

update_data_count()
root.mainloop()